package com.example.backend.service;

import lombok.Data;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.stereotype.Service;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.locks.ReentrantLock;

@Data
@Service
@ConfigurationProperties(prefix = "token-bucket")
public class TokenBucketService {
    // 用户令牌桶存储 <用户标识, 令牌桶状态>
    private final Map<String, TokenBucket> userBuckets = new ConcurrentHashMap<>();
    // 桶状态锁（按用户隔离）
    private final Map<String, ReentrantLock> bucketLocks = new ConcurrentHashMap<>();

    // 限流配置（可在配置文件中覆盖）
    private double rate = 1.0;       // 令牌生成速率 (个/秒)
    private int capacity = 5;        // 桶容量
    private long refillIntervalMs = 1000; // 令牌补充间隔(毫秒)

    // 检查是否允许请求
    public boolean allowRequest(String userKey) {
        // 添加日志
        System.out.println("Checking rate limit for user: " + userKey);

        ReentrantLock lock = bucketLocks.computeIfAbsent(userKey, k -> new ReentrantLock());
        lock.lock();
        try {
            TokenBucket bucket = userBuckets.computeIfAbsent(userKey,
                    k -> {
                        System.out.println("Creating new bucket for user: " + userKey);
                        return new TokenBucket(System.currentTimeMillis(), capacity);
                    }
            );

            boolean allowed = bucket.tryConsumeToken(System.currentTimeMillis(), rate, capacity, refillIntervalMs);

            // 添加日志
            System.out.println("User: " + userKey + " allowed: " + allowed +
                    " | Tokens: " + bucket.tokens);
            return allowed;
        } finally {
            lock.unlock();
        }
    }

    // TokenBucket
    private static class TokenBucket {
        private long lastRefillTime;  // 最后补充时间(毫秒)
        private double tokens;        // 当前令牌数

        public TokenBucket(long lastRefillTime, double tokens) {
            this.lastRefillTime = lastRefillTime;
            this.tokens = tokens;
        }

        boolean tryConsumeToken(long currentTime, double rate, int capacity, long refillInterval) {
            // 1. 精确计算应补充的令牌数（基于时间差）
            long elapsedTime = currentTime - lastRefillTime;
            if (elapsedTime > 0) {
                double tokensToAdd = ((double) elapsedTime / refillInterval) * rate;
                tokens = Math.min(tokens + tokensToAdd, capacity);
                lastRefillTime = currentTime; // 更新补充时间
            }

            // 2. 尝试消费令牌
            if (tokens >= 1) {
                tokens -= 1;
                return true;
            }
            return false;
        }
    }
}